package com.dkd.manage.task.allocator.algorithm;

import lombok.extern.slf4j.Slf4j;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
public class QLearningSelector {
    private double learningRate = 0.1;
    private double discountFactor = 0.9;
    private double explorationRate = 0.3;
    private Map<String, Double> qTable = new ConcurrentHashMap<>();

    public Long select(Map<Long, Double> scores, List<Long> candidateIds) {
        log.info("开始Q学习选择，候选人数: {}", candidateIds.size());

        if (candidateIds.isEmpty()) {
            return null;
        }

        String state = generateState(scores, candidateIds);

        // 探索还是利用
        Long selectedId;
        if (Math.random() < explorationRate) {
            // 探索：随机选择
            selectedId = candidateIds.get((int) (Math.random() * candidateIds.size()));
            log.debug("Q学习探索模式，随机选择: {}", selectedId);
        } else {
            // 利用：选择Q值最高的动作
            selectedId = exploit(state, candidateIds);
            log.debug("Q学习利用模式，选择: {}", selectedId);
        }

        // 更新Q值（简化）
        updateQValue(state, selectedId, scores.getOrDefault(selectedId, 0.0));

        log.info("Q学习选择完成，选择员工: {}", selectedId);
        return selectedId;
    }

    private String generateState(Map<Long, Double> scores, List<Long> candidateIds) {
        // 生成状态标识
        StringBuilder stateBuilder = new StringBuilder();
        for (Long empId : candidateIds) {
            double score = scores.getOrDefault(empId, 0.0);
            stateBuilder.append(empId).append(":").append((int) (score * 10)).append(";");
        }
        return stateBuilder.toString();
    }

    private Long exploit(String state, List<Long> candidateIds) {
        double maxQValue = Double.NEGATIVE_INFINITY;
        Long bestAction = candidateIds.get(0);

        for (Long action : candidateIds) {
            String key = state + "->" + action;
            double qValue = qTable.getOrDefault(key, 0.0);

            if (qValue > maxQValue) {
                maxQValue = qValue;
                bestAction = action;
            }
        }

        return bestAction;
    }

    private void updateQValue(String state, Long action, double reward) {
        String key = state + "->" + action;
        double oldQValue = qTable.getOrDefault(key, 0.0);

        // Q学习更新公式
        double newQValue = oldQValue + learningRate * (reward + discountFactor * getMaxQValue() - oldQValue);
        qTable.put(key, newQValue);

        log.debug("更新Q值: {} = {} -> {}", key, oldQValue, newQValue);
    }

    private double getMaxQValue() {
        return qTable.values().stream().max(Double::compare).orElse(0.0);
    }
}